from datetime import datetime

import akshare as ak
from matplotlib import pyplot as plt
import pandas as pd
import plotly.express as px

def rolling_annualized_return(data, window_years=3):
    """
    计算滚动年化收益率
    参数：
    data : DataFrame，必须包含date和nav两列
    window_years : 滚动窗口年数
    """
    # 预处理数据
    df = data.copy()
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date').set_index('date')

    # 计算滚动窗口天数（假设每年252个交易日）
    window_days = window_years * 252

    # 计算滚动收益率
    rolling_return = (df['close_price'].pct_change(periods=window_days) + 1)

    # 年化处理：(1 + total_return)^(1/years) - 1
    annualized_return = (rolling_return ** (1 / window_years) - 1) * 100

    return annualized_return.dropna()

def judgeNiushiData():
    # 中证偏股型基金指数(CSI:930950)
    endDate = datetime.now().strftime("%Y%m%d")
    stock_zh_index_hist_csindex_df = ak.stock_zh_index_hist_csindex(symbol="930950", start_date="20160101",
                                                                    end_date=endDate)
    print(stock_zh_index_hist_csindex_df)
    df = stock_zh_index_hist_csindex_df[['日期','指数代码','收盘']]
    df = df.rename(columns={'日期': 'date', '指数代码': "code",'收盘':'close_price'})
    df.to_csv("data/930950.csv")

def judgeNiushi():
    # 生成示例数据（2005-2023年沪深300指数）
    df = pd.read_csv("data/930950.csv")

    # 计算3年滚动年化收益
    result = rolling_annualized_return(df, window_years=3)

    # 绘制结果
    plt.figure(figsize=(12, 6))
    plt.plot(result.index, result, label='3年滚动年化收益率', color='#2c7fb8')

    # 格式设置
    plt.title('3年滚动年化收益率走势', fontsize=14, pad=20)
    plt.ylabel('年化收益率 (%)', fontsize=12)
    plt.grid(True)
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统
    # plt.rcParams['font.sans-serif'] = ['PingFang HK']  # MacOS
    # plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux
    plt.rcParams['axes.unicode_minus'] = False
    plt.legend()
    plt.tight_layout()
    plt.show()
    # 绘制图表
    fig = px.line(x=result.index, y=result, labels={'x': '日期', 'y': "收益率"},
                  title="偏股基金 3年滚动年化收益率")
    # 显示图表
    fig.show()

if __name__ == '__main__':
    judgeNiushiData()
    judgeNiushi()